"""
Author: xiash$
Date: 2025/7/20$
Project: open-image$
Description: $
***
* _ooOoo_
* o8888888o
* 88" . "88
* (| -_- |)
*  O\ = /O
* ___/`---'\____
* .   ' \\| |// `.
* / \\||| : |||// \
* / _||||| -:- |||||- \
* | | \\\ - /// | |
* | \_| ''\---/'' | |
* \ .-\__ `-` ___/-. /
* ___`. .' /--.--\ `. . __
* ."" '< `.___\_<|>_/___.' >'"".
* | | : `- \`.;`\ _ /`;.`/ - ` : | |
* \ \ `-. \_ __\ /__ _/ .-` / /
* ======`-.____`-.___\_____/___.-`____.-'======
* `=---='
*          .............................................
*           佛曰：bug泛滥，我已瘫痪！
"""

# rsio/utils/palette_converter.py
import json
import numpy as np
import logging
from osgeo import gdal, ogr, osr
from typing import Union, Type, Generator, Iterable
from collections.abc import Sequence
from .decorators import type_check, require_python_library
from .geo import safe_create_shp_field, read_image_data, copy_data_source

# 设置全局UTF-8编码
gdal.SetConfigOption("GDAL_FILENAME_IS_UTF8", "YES")
gdal.SetConfigOption("SHAPE_ENCODING", "UTF-8")
logger = logging.getLogger(__name__)


class ColorTypeConverter:
    color_map = {
        '黑色': (0, 0, 0), 'black': (0, 0, 0),
        '红色': (255, 0, 0), 'red': (255, 0, 0),
        '绿色': (0, 255, 0), 'green': (0, 255, 0),
        '蓝色': (0, 0, 255), 'blue': (0, 0, 255),
        '灰色': (128, 128, 128), 'gray': (128, 128, 128),
        '白色': (255, 255, 255), 'white': (255, 255, 255),
        '粉色': (255, 192, 203), 'pink': (255, 192, 203),
        '紫色': (128, 0, 128), 'purple': (128, 0, 128),
        # 橙色系列
        '橙色': (255, 165, 0), 'orange': (255, 165, 0),
        '浅橙色': (255, 200, 128), 'lightorange': (255, 200, 128),
        '深橙色': (255, 140, 0), 'darkorange': (255, 140, 0),
        # 黄色系列
        '黄色': (255, 255, 0), 'yellow': (255, 255, 0),
        '浅黄色': (255, 255, 224), 'lightyellow': (255, 255, 224),
        '金黄色': (255, 215, 0), 'gold': (255, 215, 0),
        # 绿色系列
        '浅绿色': (144, 238, 144), 'lightgreen': (144, 238, 144),
        '深绿色': (0, 100, 0), 'darkgreen': (0, 100, 0),
        '森林绿': (34, 139, 34), 'forestgreen': (34, 139, 34),
        '薄荷绿': (152, 251, 152), 'mintgreen': (152, 251, 152),
        '橄榄绿': (128, 128, 0), 'olive': (128, 128, 0),
        '青绿色': (0, 128, 128), 'teal': (0, 128, 128),
        # 蓝色系列
        '浅蓝色': (173, 216, 230), 'lightblue': (173, 216, 230),
        '深蓝色': (0, 0, 139), 'darkblue': (0, 0, 139),
        '天蓝色': (135, 206, 235), 'skyblue': (135, 206, 235),
        '海军蓝': (0, 0, 128), 'navy': (0, 0, 128),
        '青色': (0, 255, 255), 'cyan': (0, 255, 255),
        '宝石蓝': (65, 105, 225), 'royalblue': (65, 105, 225),
        # 紫色系列
        '浅紫色': (221, 160, 221), 'plum': (221, 160, 221),
        '深紫色': (75, 0, 130), 'indigo': (75, 0, 130),
        '薰衣草紫': (230, 230, 250), 'lavender': (230, 230, 250),
        '洋红色': (255, 0, 255), 'magenta': (255, 0, 255),
        # 棕色系列
        '棕色': (165, 42, 42), 'brown': (165, 42, 42),
        '浅棕色': (210, 180, 140), 'tan': (210, 180, 140),
        '巧克力色': (210, 105, 30), 'chocolate': (210, 105, 30),
        '沙褐色': (244, 164, 96), 'sandybrown': (244, 164, 96),
        # 粉色系列
        '浅粉色': (255, 182, 193), 'lightpink': (255, 182, 193),
        '热粉色': (255, 105, 180), 'hotpink': (255, 105, 180),
        '深粉色': (255, 20, 147), 'deeppink': (255, 20, 147),
        # 灰色系列
        '浅灰色': (211, 211, 211), 'lightgray': (211, 211, 211),
        '深灰色': (169, 169, 169), 'darkgray': (169, 169, 169),
        '银色': (192, 192, 192), 'silver': (192, 192, 192),
        '炭灰色': (105, 105, 105), 'dimgrey': (105, 105, 105),
        # 其他颜色
        '珊瑚色': (255, 127, 80), 'coral': (255, 127, 80),
        '番茄红': (255, 99, 71), 'tomato': (255, 99, 71),
        '橙红色': (255, 69, 0), 'orangered': (255, 69, 0),
        '深红色': (139, 0, 0), 'darkred': (139, 0, 0),
        '鲑鱼色': (250, 128, 114), 'salmon': (250, 128, 114),
        '海绿色': (46, 139, 87), 'seagreen': (46, 139, 87),
        '春绿色': (0, 255, 127), 'springgreen': (0, 255, 127),
        '孔雀石绿': (0, 250, 154), 'mediumspringgreen': (0, 250, 154),
        '宝石绿': (60, 179, 113), 'mediumseagreen': (60, 179, 113),
        '孔雀蓝': (51, 153, 255), 'turquoise': (64, 224, 208),
        '暗紫色': (153, 50, 204), 'darkorchid': (153, 50, 204),
        '紫罗兰': (238, 130, 238), 'violet': (238, 130, 238),
        '桃色': (255, 218, 185), 'peachpuff': (255, 218, 185),
        '米色': (245, 245, 220), 'beige': (245, 245, 220),
        '象牙色': (255, 255, 240), 'ivory': (255, 255, 240),
        '蜜色': (240, 255, 255), 'honeydew': (240, 255, 255),
        '薄雾玫瑰': (255, 228, 225), 'mistyrose': (255, 228, 225),
        '雪白色': (255, 250, 250), 'snow': (255, 250, 250),
    }

    @staticmethod
    def hex_to_rgb_tuple(hex_str: str) -> tuple[int, ...]:
        """
        将十六进制颜色代码转换为 RGB 颜色值
        :param hex_str: 十六进制颜色代码，可以是 3 位或 6 位
        :return: RGB 颜色值，格式为 (r, g, b)
        """
        hex_str = hex_str.lstrip('#')
        if len(hex_str) == 3:
            hex_str = ''.join(c * 2 for c in hex_str)  # 'F00' -> 'FF0000'
        if len(hex_str) != 6:
            raise ValueError("颜色代码必须为 6 位或 3 位十六进制数")

        return tuple(int(hex_str[i:i + 2], 16) for i in (0, 2, 4))

    @staticmethod
    def color_name_to_rgb_tuple(color_name: str) -> tuple[int, int, int]:
        """
        将颜色名称转换为 RGB 颜色值
        :param color_name: 颜色名称，例如 'red'、'blue'、'green' 等
        :return: RGB 颜色值，格式为 (r, g, b)
        """
        if (color_name_lower := color_name.lower()) not in ColorTypeConverter.color_map:
            raise ValueError(f"未知的颜色名称：{color_name}")
        return ColorTypeConverter.color_map[color_name_lower]

    @staticmethod
    def to_rgb_tuple(raw_rgb: str | list | tuple | bytes) -> tuple[int, ...]:
        match raw_rgb:
            case str() if raw_rgb.startswith('#'):
                return ColorTypeConverter.hex_to_rgb_tuple(raw_rgb)
            case str() if raw_rgb.lower() in ColorTypeConverter.color_map:
                return ColorTypeConverter.color_name_to_rgb_tuple(raw_rgb)
            case list() | tuple() if len(raw_rgb) in (3, 4) and all(isinstance(i, int) for i in raw_rgb):
                return tuple(raw_rgb[:3])
            case bytes():
                hex_str = raw_rgb.decode('utf-8').strip('#')
                return ColorTypeConverter.hex_to_rgb_tuple(hex_str)
            case _:
                raise ValueError(f"RGB值格式错误：{raw_rgb}")


class PaletteConverter:
    @staticmethod
    def convert(palette_style: str, palette: Union[list, tuple, dict, bytes, np.ndarray, gdal.ColorTable]):
        """
        根据风格类型选择转换策略
        Args:
            palette_style: 'pil'/'gdal'/'rasterio'。。。
            palette: 任意格式的调色板
        Returns:
            后端所需的格式
        """
        match palette_style:
            case 'pil':
                return PaletteConverter._to_pil(palette)
            case 'gdal':
                return PaletteConverter._to_gdal(palette)
            case 'rasterio':
                return PaletteConverter._to_rasterio(palette)
            case _:
                raise ValueError(f'不支持的backend_type: {palette_style}')

    # PIL格式：[r,g,b,r,g,b...]
    @staticmethod
    def _to_pil(
            palette: Union[
                dict[int, tuple | list | str | bytes],
                list[tuple | list],
                tuple[list | tuple | str | bytes],
                bytes,
                np.ndarray
            ]
    ) -> list[int]:
        # 特殊处理 gdal.ColorTable 类型
        if isinstance(palette, gdal.ColorTable):
            return PaletteConverter.gdal_color_table_to_pil(palette)
        match palette:
            case dict() if all(
                    isinstance(k, int) and isinstance(v, Sequence) for k, v in palette.items()):
                palette_pil = []
                for class_id in range(max(palette) + 1):
                    palette_pil.extend(ColorTypeConverter.to_rgb_tuple(palette.get(class_id, (0, 0, 0))))
                return palette_pil
            case bytes():
                return PaletteConverter._bytes_to_rgb_pil(palette)
            case np.ndarray():
                if palette.ndim == 1:
                    return PaletteConverter._validate_pil_palette(palette.tolist())
                elif palette.ndim == 2 and palette.shape[1] == 3:
                    return PaletteConverter._validate_pil_palette(palette.flatten().tolist())
                else:
                    raise ValueError(f"调色板格式不支持的 numpy.ndarray 维度: {palette.shape}")
            case list() | tuple() if all(isinstance(x, int) for x in palette):
                return PaletteConverter._validate_pil_palette(palette)
            case list() | tuple() if all(isinstance(x, Sequence) for x in palette):
                # Sequence[Sequence[int]] 且每个子序列是长度为 3 的整数列表
                return PaletteConverter._validate_pil_palette(
                    [v for c in palette for v in ColorTypeConverter.to_rgb_tuple(c)])
            case _:
                raise ValueError(f"不支持的调色板格式: {palette}")

    # GDAL格式：ColorTable
    @staticmethod
    def _to_gdal(palette: Union[list, dict, str, bytes, np.ndarray]) -> gdal.ColorTable:
        if isinstance(palette, gdal.ColorTable):
            return palette
        pil_palette = PaletteConverter._to_pil(palette)
        ct = gdal.ColorTable()
        for i in range(0, len(pil_palette), 3):
            ct.SetColorEntry(i // 3, tuple(pil_palette[i:i + 3]))
        return ct

    # Rasterio格式：{index: (r,g,b)}
    @staticmethod
    def _to_rasterio(palette: Union[list, tuple, bytes, dict, np.ndarray]) -> dict:
        pil_palette = PaletteConverter._to_pil(palette)
        return {i // 3: tuple(pil_palette[i:i + 3]) for i in range(0, len(pil_palette), 3)}

    # 辅助方法
    @staticmethod
    def _bytes_to_rgb_pil(palette: bytes) -> list[int]:
        """
         纯 RGB 字节流（长度为 3 的倍数）转换成pil格式的list[int,...]格式
        """
        if len(palette) % 3 != 0:
            raise ValueError("bytes 长度必须是 3 的倍数")
        return [int(v) for v in palette]

    @staticmethod
    def _validate_pil_palette(palette: Sequence[int]) -> list:
        """
        验证并返回标准PIL调色板格式
        PIL调色板要求长度为3的倍数，至少为3，且不超过768(256*3)，元素值在0-255之间
        """
        if len(palette) == 0 or len(palette) % 3 != 0:
            raise ValueError(f"PIL调色板长度必须为3的倍数且不能为0！，实际长度为：{len(palette)}")

        if not all(isinstance(x, int) and 0 <= x <= 255 for x in palette):
            raise ValueError("所有调色板值必须是0-255之间的整数")

        return list(palette)

    @staticmethod
    def gdal_color_table_to_pil(gdal_ct: gdal.ColorTable) -> list[int]:
        """
        将 gdal.ColorTable 转换为 PIL 格式的调色板。

        Args:
            gdal_ct: gdal.ColorTable 对象。

        Returns:
            list[int]: PIL 格式的调色板（[r,g,b,r,g,b,...]）。
        """
        pil_palette = []
        for i in range((min(256, gdal_ct.GetCount()))):
            entry = gdal_ct.GetColorEntry(i)
            if entry:
                pil_palette.extend(entry[:3])  # 只取 RGB 值，忽略 alpha (如果存在)
        return pil_palette

    @staticmethod
    def clone_color_table(original_ct: gdal.ColorTable) -> gdal.ColorTable:
        """
        克隆 GDAL ColorTable 对象以避免引用问题
        """
        if not original_ct:
            return original_ct
        new_ct = gdal.ColorTable()
        # GDAL ColorTable 最多支持 256 个条目
        for i in range(min(256, original_ct.GetCount())):
            entry = original_ct.GetColorEntry(i)
            if entry:
                new_ct.SetColorEntry(i, entry)
        return new_ct


class ArrayShapeConverter:
    """
    array数组形状转换器
    """

    @staticmethod
    def convert(shape_style: str, array: np.ndarray) -> np.ndarray:
        """
        根据指定风格转换数组形状
        Args:
            shape_style: 目标形状风格 ('gdal', 'rasterio', 'pil', 'cv2')
            array: 输入的numpy数组
        Returns:
            转换后的numpy数组
        Raises:
            ValueError: 不支持的shape_style或数组维度不正确
        """

        match shape_style:
            case 'gdal' | 'rasterio':
                array = ArrayShapeConverter.array_shape_to_gdal(array)
            case 'pil' | 'cv2':
                array = ArrayShapeConverter.array_shape_to_cv2(array)
            case _:
                raise ValueError(f"暂不支持的shape_style：{shape_style}")
        return array

    @staticmethod
    @type_check
    def array_shape_to_cv2(array: np.ndarray) -> np.ndarray:
        """
        将数组的形状转换为cv2的格式
        :param array: numpy数组
        :return: 转换后的数组和是否需要交换通道
        """
        if array.ndim not in (2, 3):
            raise ValueError("数组的维度必须为2或3！")
        if array.ndim == 2:
            return array
        if np.argmin(array.shape) == 0:
            array = array.transpose(1, 2, 0)
        return array

    @staticmethod
    def array_shape_to_pil(array: np.ndarray) -> np.ndarray:
        return ArrayShapeConverter.array_shape_to_cv2(array)

    @staticmethod
    @type_check
    def array_shape_to_gdal(array: np.ndarray) -> np.ndarray:
        if array.ndim not in (2, 3):
            raise ValueError("数组的维度必须为2或3！")
        if array.ndim == 2:
            return array
        if np.argmin(array.shape) == 2:
            array = array.transpose((2, 0, 1))
        return array

    @staticmethod
    def array_shape_to_rasterio(array: np.ndarray) -> np.ndarray:
        array = ArrayShapeConverter.array_shape_to_gdal(array)
        if len(array.shape) == 2:
            array = array[np.newaxis, ...]
        return array

    @staticmethod
    def array_shape_to_rgb(array: np.ndarray) -> np.ndarray:
        array = ArrayShapeConverter.array_shape_to_rasterio(array)
        if array.shape[0] == 1:
            array = np.tile(array, (3, 1, 1))
        elif array.shape[0] > 3:
            array = array[2::-1, ...]  # 超过3波段的图像rgb一般是反过来的
        else:
            logger.debug('图像波段数不满足条件，转换rgb失败！')
        return array


class DataTypeConverter:
    """
    数据类型转换器
    """

    @staticmethod
    def np_type_to_gdal(dtype: np.dtype | Type[np.number]) -> int:
        """
        将 NumPy 数据类型映射到 GDAL 数据类型常量。

        Args:
            dtype (np.dtype): NumPy 数据类型。

        Returns:
            int: GDAL 数据类型常量。
        """
        np_gdal_types = {
            np.uint8: gdal.GDT_Byte,
            np.int8: gdal.GDT_Int16,
            np.uint16: gdal.GDT_UInt16,
            np.int16: gdal.GDT_Int16,
            np.uint32: gdal.GDT_UInt32,
            np.int32: gdal.GDT_Int32,
            np.float32: gdal.GDT_Float32,
            np.float64: gdal.GDT_Float64,
        }
        if isinstance(dtype, np.dtype):
            dtype = dtype.type
        return np_gdal_types.get(dtype, gdal.GDT_Unknown)

    @staticmethod
    def gdal_type_to_np(gdal_type: int) -> Type[np.number]:
        """
        将 GDAL 数据类型映射到 NumPy 数据类型。

        Args:
            gdal_type (int): GDAL 数据类型常量。

        Returns:
            np.dtype: NumPy 数据类型。
        """
        gdal_np_types = {
            gdal.GDT_Byte: np.uint8,
            gdal.GDT_Int16: np.int16,
            gdal.GDT_UInt16: np.uint16,
            gdal.GDT_Int32: np.int32,
            gdal.GDT_UInt32: np.uint32,
            gdal.GDT_Float32: np.float32,
            gdal.GDT_Float64: np.float64,
        }
        return gdal_np_types.get(gdal_type, np.uint8)

    @staticmethod
    def type_to_shp_field(data_type):
        shp_field_types = {
            "int": ogr.OFTInteger,
            "float": ogr.OFSTFloat32,
            "str": ogr.OFTString,
            "string": ogr.OFTString,
            int: ogr.OFTInteger,
            float: ogr.OFSTFloat32,
            str: ogr.OFTString,
            ogr.OFTInteger: ogr.OFTInteger,
            ogr.OFTString: ogr.OFTString,
            ogr.OFSTFloat32: ogr.OFSTFloat32
        }
        dst_type = shp_field_types.get(data_type)
        if dst_type is None:
            raise ValueError(f"不支持的shp数据类型：{data_type}")
        return dst_type


# 坐标类型转换
class GeoConverter:
    """
    坐标转换器
    """

    @staticmethod
    def pixel_to_geo(x: int, y: int, geo_transform: tuple[float, ...]) -> tuple[float, float]:
        """
        将像素坐标转换为地理坐标
        :param x: 像素坐标x
        :param y: 像素坐标y
        :param geo_transform: 地理坐标转换参数
        :return: 地理坐标(lon, lat)
        """
        lon = geo_transform[0] + x * geo_transform[1] + y * geo_transform[2]
        lat = geo_transform[3] + x * geo_transform[4] + y * geo_transform[5]
        return lon, lat

    @staticmethod
    def geo_to_pixel(lon: float, lat: float, geo_transform: tuple[float, ...]) -> tuple[int, int]:
        """
        将地理坐标转换为像素坐标
        :param lon: 地理坐标x
        :param lat: 地理坐标y
        :param geo_transform: 地理坐标转换参数
        :return: 像素坐标(x, y)
        """
        x = int((lon - geo_transform[0]) / geo_transform[1])
        y = int((lat - geo_transform[3]) / geo_transform[5])
        return x, y

    @staticmethod
    @require_python_library('affine', 'affine')
    def geo_transform2rasterio(geo_transform: tuple[float, ...]) -> 'affine.Affine':
        """
        转换transform为affine格式
        :param geo_transform:
        :return:
        """
        from affine import Affine
        if isinstance(geo_transform, Affine):
            return geo_transform
        if isinstance(geo_transform, tuple) and len(geo_transform) == 6:  # 当地理坐标仿射变换是gdal格式时，转换为affine格式
            return Affine.from_gdal(*geo_transform)
        logger.error('[red]geo_transform格式出错！只支持(x,dx,rx,y,ry,dy)转换affine格式！', extra={'markup': True})

    @staticmethod
    def geo_transform2gdal(geo_transform: 'affine.Affine') -> tuple[float, ...]:
        """
        转换affine格式为gdal格式
        :param geo_transform:
        :return:
        """
        from affine import Affine
        if isinstance(geo_transform, Affine):
            return geo_transform.to_gdal()
        if isinstance(geo_transform, tuple) and len(geo_transform) == 6:
            return geo_transform
        logger.warning('[yellow]geo_transform格式错误！只支持Affine格式转gdal格式', extra={'markup': True})

    @staticmethod
    def pixel_bounds2geo_bounds(
            pixel_bounds: tuple[int, int, int, int],
            geo_transform: tuple[float, float, float, float, float, float]
    ) -> tuple[float, float, float, float]:
        """
        将像素边界转换为地理边界
        :param pixel_bounds:
        :param geo_transform:
        :return:
        """
        pixel_left, pixel_top, pixel_right, pixel_bottom = pixel_bounds
        geo_left, geo_top = GeoConverter.pixel_to_geo(pixel_left, pixel_top, geo_transform)
        geo_right, geo_bottom = GeoConverter.pixel_to_geo(pixel_right, pixel_bottom, geo_transform)
        return geo_left, geo_top, geo_right, geo_bottom

    @staticmethod
    def geo_bounds2pixel_bounds(
            geo_bounds: tuple[float, float, float, float],
            geo_transform: tuple[float, float, float, float, float, float]
    ) -> tuple[int, int, int, int]:
        geo_left, geo_top, geo_right, geo_bottom = geo_bounds
        pixel_left, pixel_top = GeoConverter.geo_to_pixel(geo_left, geo_top, geo_transform)
        pixel_right, pixel_bottom = GeoConverter.geo_to_pixel(geo_right, geo_bottom, geo_transform)
        return pixel_left, pixel_top, pixel_right, pixel_bottom

    @staticmethod
    def shp_to_raster(
            src_ds: str | ogr.DataSource,
            out_width: int,
            out_height: int,
            output_path: str = '',
            attribute_field: str = 'class_id',
            mapping_field: dict = None,
            field_to_create: str = 'class_id',
            geo_transform: tuple = None,
            projection: str = None,
            nodata_value: float = 0,
            data_type: int = gdal.GDT_Byte
    ) -> gdal.Dataset:
        """
        将矢量文件转换为栅格文件
        Args:
            src_ds: 输入的矢量文件路径或者ogr.DataSource对象
            output_path: 输出的栅格文件路径
            attribute_field: 用于栅格化的属性字段名（可选）
            mapping_field: 字段映射，例如：{'建筑':1,'道路':2}
            field_to_create: 有字段映射时，自动创建的新字段名称
            out_width: 输出栅格的列数
            out_height: 输出栅格的行数
            geo_transform: 地理变换参数
            projection: 投影信息
            nodata_value: NoData值
            data_type: 数据类型
        """
        if isinstance(src_ds, str):
            # 打开矢量文件
            src_ds = ogr.Open(src_ds, 1)
            if src_ds is None:
                raise ValueError(f"无法打开矢量文件: {src_ds}")
        layer = src_ds.GetLayer()
        # 如果没有提供地理变换参数，则从矢量数据获取
        if geo_transform is None or projection is None:
            extent = layer.GetExtent()
            x_min, x_max, y_min, y_max = extent
            # 计算像素大小
            x_res = (x_max - x_min) / out_width
            y_res = (y_max - y_min) / out_height
            geo_transform = (x_min, x_res, 0, y_max, 0, -y_res)
            projection = layer.GetSpatialRef().ExportToWkt()
        if mapping_field and attribute_field:
            safe_create_shp_field(layer, field_to_create, ogr.OFTInteger)
            # 更新每个要素
            for feature in layer:
                class_name = feature.GetField(attribute_field)
                burn_value = mapping_field.get(class_name, 0)
                feature.SetField(field_to_create, burn_value)
                layer.SetFeature(feature)
            attribute_field = field_to_create
        # 创建输出栅格数据集
        driver = gdal.GetDriverByName('MEM' if output_path == '' else 'GTiff')
        raster_ds = driver.Create(
            output_path,
            out_width,
            out_height,
            1,  # 波段数
            data_type
        )
        raster_ds.SetGeoTransform(geo_transform)
        raster_ds.SetProjection(projection)
        # 设置NoData值
        band = raster_ds.GetRasterBand(1)
        band.SetNoDataValue(nodata_value)
        band.Fill(nodata_value)
        # 栅格化选项
        if attribute_field:
            options = [f"ATTRIBUTE={attribute_field}"]
        else:
            options = [f"BURN_VALUE=1"]
        # 执行栅格化
        gdal.RasterizeLayer(
            raster_ds,
            [1],  # 波段列表
            layer,
            options=options
        )
        if output_path:
            raster_ds = None
        return raster_ds

    @staticmethod
    def shp_to_raster_with_template(
            shapefile_path: str,
            template_raster: str | gdal.Dataset,
            output_path: str = '',
            attribute_field: str = 'class_id',
            mapping_field: dict = None
    ):
        """
        基于现有栅格模板将矢量转换为栅格
        Args:
            shapefile_path: 输入的矢量文件路径
            template_raster: 用作模板的栅格文件路径或gdal.Dataset对象
            output_path: 输出的栅格文件路径
            attribute_field: 用于栅格化的属性字段名
            mapping_field: 字段映射，例如：{'建筑':1,'道路':2}
        """
        # 打开模板栅格
        if isinstance(template_raster, str):
            template_raster = gdal.Open(template_raster)
            if template_raster is None:
                raise ValueError(f"无法打开模板栅格文件: {template_raster}")
        assert isinstance(template_raster, gdal.Dataset), "模板栅格必须是文件路径或gdal.Dataset对象！"
        # 获取模板栅格的参数
        x_pixels = template_raster.RasterXSize
        y_pixels = template_raster.RasterYSize
        geo_transform = template_raster.GetGeoTransform()
        projection = template_raster.GetProjection()
        data_type = template_raster.GetRasterBand(1).DataType
        # 调用基本转换函数
        return GeoConverter.shp_to_raster(
            shapefile_path,
            out_width=x_pixels,
            out_height=y_pixels,
            output_path=output_path,
            attribute_field=attribute_field,
            mapping_field=mapping_field,
            geo_transform=geo_transform,
            projection=projection,
            data_type=data_type
        )

    @staticmethod
    def geo_projection2wkt(geo_projection: str | int | osr.SpatialReference):
        sr = osr.SpatialReference()
        match geo_projection:
            case int():
                sr.ImportFromEPSG(geo_projection)
            case str() if geo_projection.upper().startswith('EPSG:'):
                sr.ImportFromEPSG(int(geo_projection[5:]))
            case str():
                sr.ImportFromWkt(geo_projection)
            case osr.SpatialReference():
                sr = geo_projection
            case _:
                raise ValueError(f"不支持的投影信息: {geo_projection}")
        return sr.ExportToWkt()

    @staticmethod
    def mask_array_to_shp(
            mask_array: np.ndarray,
            save_shp_path: str = None,
            class_info: list | tuple | dict = None,
            field_info: list | tuple = (('class_id', int), ('class_name', str)),
            geo_transform: 'tuple|affine.Affine' = (0, 1, 0, 0, 0, -1),
            geo_projection: str = None,
            del_zero_value: bool = True,
            simplify_tolerance: float = 2,
    ) -> ogr.DataSource:
        """
        mask预测图array转矢量shp
        :param mask_array: mask图片路径或者array
        :param save_shp_path: 矢量shp文件保存路径
        :param class_info: 分类信息，格式同dict_model配置文件\
        :param field_info: shp字段名称及类型
        :param geo_transform: 地理坐标变换系数
        :param geo_projection: 投影信息
        :param del_zero_value: 是否删除背景0
        :param simplify_tolerance: 边缘点精简，容忍度参数，值越大精简越狠
        :return:
        """
        from rasterio import features
        assert mask_array.ndim == 2, 'mask_array必须为单通道灰度图！'
        geo_transform = GeoConverter.geo_transform2rasterio(geo_transform)
        if geo_transform.e == 1:
            from affine import Affine
            geo_transform = Affine(1, 0, 0, 0, -1, 0)
        shapes = features.shapes(mask_array, mask=(mask_array != 0) if del_zero_value else None,
                                 transform=geo_transform)  # 栅格矢量化
        shp_driver: ogr.Driver = ogr.GetDriverByName("Memory")  # 创建矢量内存驱动
        dst_shp_ds: ogr.DataSource = shp_driver.CreateDataSource('')  # 创建文件资源
        proj = osr.SpatialReference()  # 创建坐标系类
        proj.ImportFromWkt(geo_projection if geo_projection else '')  # 导入坐标系
        dst_layer: ogr.Layer = dst_shp_ds.CreateLayer('layer', geom_type=ogr.wkbPolygon, srs=proj)  # 创建多边形图层，并写入坐标系
        # 创建字段
        for field_name, field_type in field_info:
            safe_create_shp_field(dst_layer, field_name, DataTypeConverter.type_to_shp_field(field_type))
        for feature in GeoConverter.shapes_to_features(shapes, dst_layer, class_info, field_info,
                                                       simplify_tolerance * geo_transform.a):
            dst_layer.CreateFeature(feature)  # 将feature要素加到图层中
        if save_shp_path is not None:
            copy_data_source(dst_shp_ds, save_shp_path)
        return dst_shp_ds

    @staticmethod
    def shapes_to_features(
            shapes: Iterable[tuple[dict, int]],
            layer_defn: ogr.FeatureDefn | ogr.Layer,
            class_info: list | tuple | dict,
            field_info: list | tuple = (('class_id', int), ('class_name', float)),
            simplify_tolerance: float = None
    ) -> Generator[ogr.Feature, None, None]:
        """
        将rast
        @param shapes: 由rasterio的features.shapes栅格矢量化的生成器对象
        @param layer_defn: shp图层的定义信息，ogr.FeatureDefn | ogr.Layer
        @param class_info: 类别信息，如['背景','建筑',。。。],{1:'建筑',2:'道路'},{1:['建筑','高层建筑'],2:['道路','农村道路']}
        @param field_info: shp字段名称及数据类型，默认为(('class_id', int), ('class_name', float))
        @param simplify_tolerance: 矢量边缘点精简的阈值参数
        @return:
        """
        if isinstance(layer_defn, ogr.Layer):
            layer_defn = layer_defn.GetLayerDefn()
        match class_info:
            case list() | tuple():
                class_info = {i: [value] for i, value in enumerate(class_info)}
            case dict():
                class_info = {i: [value] for i, value in class_info.items() if not isinstance(value, (list, tuple))}
            case _:
                raise ValueError(f"class_info 类别信息格式错误！请查看注释说明")
        for polygon, class_id in shapes:  # 遍历shapes，每个shape中保存的是要素的类型以及坐标（字典格式），还有id值
            class_id = int(class_id)
            if class_id not in class_info:
                continue
            # 先将字典转为json格式，再从geojson格式创建geometry
            geometry: ogr.Geometry = ogr.CreateGeometryFromJson(json.dumps(polygon))
            # 简化边缘点
            if simplify_tolerance:
                geometry = geometry.Simplify(simplify_tolerance)
            feature: ogr.Feature = ogr.Feature(layer_defn)  # 创建feature要素
            field_values = [class_id] + class_info[class_id]  # 加上第一个属性字段class_id的值
            for field, field_value in zip(field_info, field_values):  # 遍历所有属性字段，逐个设置值
                feature.SetField(field[0], field_value)  # 设置字段的值
            feature.SetGeometry(geometry)  # 给feature要素添加多边形的坐标点位置信息
            yield feature

    @staticmethod
    def raster_to_shp(
            raster_path: str,
            save_shp_path: str = None,
            ref_raster_path: str = None,
            class_info: list | tuple | dict | None = None,
            field_info: list | tuple = (('class_id', int), ('class_name', str)),
            del_zero_value: bool = True,
            simplify_tolerance: float = 2
    ) -> ogr.DataSource:
        """
        栅格文件转矢量
        :param raster_path: 栅格路径
        :param save_shp_path: 矢量保存路径
        :param ref_raster_path: 用于坐标及投影的参考栅格文件路径
        :param class_info: 类别信息
        :param field_info: shp字段名称及数据类型
        :param del_zero_value: 是否删除背景值0
        :param simplify_tolerance: 矢量精简阈值
        :return:
        """
        img_data = read_image_data(raster_path)
        mask_array = img_data.ReadAsArray()
        if mask_array.ndim != 2:
            raise ValueError("图像必须为单通道灰度图！")
        if ref_raster_path is not None:
            ref_img_data = read_image_data(ref_raster_path)
            geo_transform = ref_img_data.GetGeoTransform()
            geo_projection = ref_img_data.GetProjection()
        else:
            geo_transform = img_data.GetGeoTransform()
            geo_projection = img_data.GetProjection()
        shp_data = GeoConverter.mask_array_to_shp(
            mask_array,
            save_shp_path,
            class_info=class_info,
            field_info=field_info,
            geo_transform=geo_transform,
            geo_projection=geo_projection,
            del_zero_value=del_zero_value,
            simplify_tolerance=simplify_tolerance
        )
        return shp_data
